# -*- coding: utf-8 -*-

import json
import os

from aws_sqs import Sqs


class S3Sqs(Sqs):

    def _init(self):
        super(S3Sqs, self)._init()

        self.__s3 = self.get_resource_by_name('s3')

    def send(self, data):
        if self._cfg.get('send_s3_enabled', True) is False:
            return super(S3Sqs, self).send(data)
        else:
            return super(S3Sqs, self).send(data, before_send=self.before_send)

    def receive(self):
        if self._cfg.get('receive_s3_enabled', True) is False:
            return super(S3Sqs, self).receive()
        else:
            return super(S3Sqs, self).receive(after_receive=self.after_receive)

    def after_receive(self, data):
        msg = None

        try:
            data = json.loads(data)
            records = data.get('Records', None)
            if records:
                s3 = records[0].get('s3', None)
                if s3:
                    bucket_name = s3['bucket']['name']
                    object_key = s3['object']['key']
                    fn = '%s/%s_%s' % (self._container.path('s3sqs-receive'), bucket_name, object_key.replace('/', '-'))
                    with open(fn, 'wb') as f:
                        self.__s3.Bucket(bucket_name).download_fileobj(object_key, f)
                    with open(fn, 'rb') as f:
                        msg = json.load(f)
                    os.remove(fn)
        except Exception as e:
            raise Exception('S3Sqs after_receive: [data] %s [error] %s' % (data, e))

        return msg

    def before_send(self, data):
        try:
            fn = '%s/%s_%s' % (self._container.path('s3sqs-send'), data['bucket_name'], data['object_key'].replace('/', '-'))
            with open(fn, 'wb') as f:
                f.write(data['data'])
                f.write('\r\n')
            self.__s3.Bucket(data['bucket_name']).upload_file(fn, data['object_key'])
            os.remove(fn)
        except Exception as e:
            raise Exception('S3Sqs before_send: [data] %s [error] %s' % (data, e))

        return None
